# -*- coding: utf-8 -*-

# @File    : robot_function.py
# @Date    : 2019-09-03
# @Author  : gongsunmolan
from wxpy import *
import pymysql
import math
from PIL import Image
import os
import jieba
from gensim import corpora, models, similarities

##获取好友消息并写入数据库
def get_friends(myFriends):
    print(myFriends[1].raw)
    db = pymysql.connect("localhost", "root", "", "wechat")
    cursor = db.cursor()
    cursor.execute('TRUNCATE TABLE friends')
    db.commit()
    db.close()

    for i in range(len(myFriends)):
        name = myFriends[i].name
        nick_name = myFriends[i].nick_name
        puid = myFriends[i].puid
        Province = myFriends[i].raw['Province']
        City = myFriends[i].raw['City']
        Sex = myFriends[i].raw['Sex']
        Signature = myFriends[i].raw['Signature']
        HeadImgUrl = myFriends[i].raw['HeadImgUrl']
        UserName = myFriends[i].raw['UserName']
        # 打开数据库连接
        db = pymysql.connect("localhost", "root", "", "wechat")
        # 使用cursor()方法获取操作游标
        cursor = db.cursor()
        # SQL 插入语句
        sql = """INSERT INTO friends(name,nick_name,puid,Province,City,Sex,Signature,HeadImgUrl,UserName)VALUES ('%s','%s','%s','%s','%s','%s','%s','%s','%s')""" % (
        name, nick_name, puid, Province, City, Sex, Signature, HeadImgUrl, UserName)
        try:
            # 执行sql语句
            cursor.execute(sql)
            # 提交到数据库执行
            db.commit()
        except:
            # 如果发生错误则回滚
            db.rollback()
        # 关闭数据库连接
        db.close()


# 清空文件夹文件
def del_file(path):
    ls = os.listdir(path)
    for i in ls:
        c_path = os.path.join(path, i)
        if os.path.isdir(c_path):
            del_file(c_path)
        else:
            os.remove(c_path)


# 获取好友头像并拼接
def get_friends_heardimg(bot, msg):
    # 获取当前路径信息
    curr_dir = os.getcwd()
    print(curr_dir)

    print("头像处理")
    # avater_dir = os.path.join(os.getcwd(),'wechat\\1\\')
    if not os.path.exists(curr_dir + "/Resources/FriendImgs/"):
        os.mkdir(curr_dir + "/Resources/FriendImgs/")

    avatar_dir = curr_dir + "/Resources/FriendImgs/"
    del_file(avatar_dir)

    friends = bot.friends(update=True)
    num = 0
    for friend in friends:
        friend.get_avatar(avatar_dir + '\\' + str(num) + ".jpg")
        print('好友昵称:%s' % friend.nick_name)
        num = num + 1
    # 获取文件内头像个数
    length = len(os.listdir(avatar_dir))
    # 设置画布大小
    image_size = 2500
    # 设置每个头像大小
    each_size = math.ceil(2560 / math.floor(math.sqrt(length)))
    # 计算所需各行列的头像数量
    x_lines = math.ceil(math.sqrt(length))
    y_lines = math.ceil(math.sqrt(length))
    image = Image.new('RGB', (each_size * x_lines, each_size * y_lines))
    x = 0
    y = 0
    for (root, dirs, files) in os.walk(avatar_dir):
        for pic_name in files:
            # 增加头像读不出来的异常处理
            try:
                with Image.open(avatar_dir + pic_name) as img:
                    img = img.resize((each_size, each_size))
                    image.paste(img, (x * each_size, y * each_size))
                    x += 1
                    if (x == x_lines):
                        x = 0
                        y += 1
            except IOError:
                print("头像读取失败")
    img = image.save(os.getcwd() + "/wechat.png")
    # bot.file_helper.send_image(avatar_dir + "/wechat.png")
    print('微信好友头像拼接完成!')
    return 1


# 关键词回复
def ireply(bot, msg):
    print("回复函数内部")
    keywords = msg.text
    print(type(keywords))

    # 打开数据库连接
    db = pymysql.connect("localhost", "root", "", "wechat")
    # 使用cursor()方法获取操作游标
    cursor = db.cursor()
    # SQL 插入语句
    # sql = """SELECT * FROM  knowledge keywords =="""" +  keywords +"/""

    sql = """select content from knowledge where keywords = "%s" """ % (keywords)
    print(sql)
    try:
        # 执行sql语句
        cursor.execute(sql)
        # 获取所有记录列表
        results = cursor.fetchall()
        for row in results:  # 从fetchall中读取操作

            print(row)
            str1 = str(row)
            print(str1)
            str2 = str1[2:-3]
            print(str2)

    except:
        # 如果发生错误则回滚
        db.rollback()
    # 关闭数据库连接
    db.close()
    if str2 is not None:
        return str2


def similre(bot, msg):
    print("相似度回复函数内部")
    doc_test = msg.text
    print(type(doc_test))

    # 添加自定义词典
    # file_name 为文件类对象或自定义词典的路径
    jieba.load_userdict('E:/03/wechat_robot/Resources/dic/dict.txt')

    all_doc = []

    # 打开数据库连接
    db = pymysql.connect("localhost", "root", "", "wechat")
    # 使用cursor()方法获取操作游标
    cursor = db.cursor()
    # SQL 插入语句
    # sql = """SELECT * FROM  knowledge keywords =="""" +  keywords +"/""

    sql = """select keywords from knowledge """
    print(sql)
    try:
        # 执行sql语句
        cursor.execute(sql)
        # 获取所有记录列表
        results = cursor.fetchall()
        for row in results:  # 从fetchall中读取操作

            # print(row)
            str1 = str(row)
            # print(str1)
            str2 = str1[2:-3]
            # print(str2)
            all_doc.append(str2)

    except:
        # 如果发生错误则回滚
        db.rollback()
    # 关闭数据库连接
    db.close()

    all_doc_list = []
    for doc in all_doc:
        doc_list = [word for word in jieba.cut(doc)]
        all_doc_list.append(doc_list)

    print(all_doc_list)

    doc_test_list = [word for word in jieba.cut(doc_test)]
    print(doc_test_list)

    dictionary = corpora.Dictionary(all_doc_list)
    dictionary.keys()
    dictionary.token2id
    corpus = [dictionary.doc2bow(doc) for doc in all_doc_list]

    doc_test_vec = dictionary.doc2bow(doc_test_list)
    doc_test_vec
    print(doc_test_vec)
    tfidf = models.TfidfModel(corpus)
    tfidf[doc_test_vec]

    index = similarities.SparseMatrixSimilarity(tfidf[corpus], num_features=len(dictionary.keys()))
    sim = index[tfidf[doc_test_vec]]

    # print(sim)
    a = sorted(enumerate(sim), key=lambda item: -item[1])

    print(a[0][0])
    print(a[0][1])
    # for x in a:
    #    print(x)

    if a[0][1] > 0.5:
        print("相似度满足条件")
        print(all_doc[a[0][0]])
        keywords = all_doc[a[0][0]]
        db = pymysql.connect("localhost", "root", "", "wechat")
        cursor = db.cursor()
        sql = """select content from knowledge where keywords = "%s" """ % (keywords)
        print(sql)
        try:
            cursor.execute(sql)
            results = cursor.fetchall()
            for row in results:  # 从fetchall中读取操作
                # print(row)
                str1 = str(row)
                # print(str1)
                str2 = str1[2:-3]
                # print(str2)
        except:
            # 如果发生错误则回滚
            db.rollback()
        # 关闭数据库连接
        db.close()
        if str2 is not None:
            return str2
